import asyncio,telnetlib3

import getopt, sys
import datetime
from pathlib import Path
from sigmf import SigMFFile
import socket

outDir = "."
IPaddr  = "192.168.1.99"
centerFreq = 1000.0
port = 13000
init = False
setip = False
sigmf_d = None

def get_ip():
    hostname = socket.gethostname()
    ip_address = socket.gethostbyname(hostname)
    if "192.168.1" in ip_address:
        IPaddr = ip_address
        return True
    else:
        print("bad ip address",ip_address)
        return False
    
def data_server():
    now = datetime.datetime.utcnow()
    formatted_date = now.strftime("%Y_%m_%d_%H_%M_%S")
    print(f"transmitting {formatted_date}")
    dat_fn = f'rx_{formatted_date}_{centerFreq}m.sigmf-data'
    dat_dir_fn = Path(outDir).joinpath(dat_fn)
    sigmf_d = dat_dir_fn.open('wb')
    
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    server_socket.bind(('0.0.0.0', port))

    server_socket.listen(5)

    print("TCP server start, wait client...")

    total = 0
    end = False
    while not end:

        client_socket, client_address = server_socket.accept()
    
        print("connect with {}".format(client_address))
    
        while total<256*1024*1024:

            data = client_socket.recv(8192)
            if data:
                print("receive {}: {}".format(client_address, len(data)))
            
                sigmf_d.write(data)
            else:
                sigmf_d.close()
                print(f"end {total}")
                end = True
                break
        print("prepare meta")
        meta = SigMFFile(
            data_file=str(dat_dir_fn), # extension is optional
            global_info = {
                SigMFFile.DATATYPE_KEY: 'ci16_le',  # in this case, 'ci16_le'
                SigMFFile.SAMPLE_RATE_KEY: 122880000,
                SigMFFile.AUTHOR_KEY: 'a4a881d4',
                SigMFFile.DESCRIPTION_KEY: f'{centerFreq}M',
                SigMFFile.VERSION_KEY : '1.0.0',
            }
        )
        freq = int(round(centerFreq*1e6))
        meta.add_capture(0, metadata={
            SigMFFile.FREQUENCY_KEY: freq,
            SigMFFile.DATETIME_KEY: now.isoformat()+'Z',
        })
        mete_fn = f'rx_{formatted_date}_{centerFreq}m.sigmf-meta'
        meta.tofile(str(Path(outDir).joinpath(mete_fn)))
        print("client closed")
        client_socket.close()
    print("server closed")
    server_socket.close()
    
def usage():
    print("Usage: python capture.py [-i] [-f CentFreq] [-d Dir] [-p port]")
    
async def handle_client(reader, writer):
    now = datetime.datetime.now()
    formatted_date = now.strftime("%Y_%m_%d_%H_%M_%S")
    print(f"transmitting {formatted_date}")
    dat_fn = f'rx_{formatted_date}_{centerFreq}m.sigmf-data'
    sigmf_d = Path(outDir).joinpath(dat_fn).open('wb')
    while True:
        data = await reader.read(8192)
        if data:
            len = len(data)
            print(f"Received {len}",flush=True)
        
            sigmf_d.write(data)
        else:
            sigmf_d.close()
            break
    print("prepare meta")
    meta = SigMFFile(
        data_file=dat_fn, # extension is optional
        global_info = {
            SigMFFile.DATATYPE_KEY: 'ci16_le',  # in this case, 'ci16_le'
            SigMFFile.SAMPLE_RATE_KEY: 122880000,
            SigMFFile.AUTHOR_KEY: 'a4a881d4',
            SigMFFile.DESCRIPTION_KEY: f'{centerFreq}M',
            SigMFFile.VERSION_KEY : '1.0.0',
        }
    )
    freq = int(round(centerFreq*1e6))
    meta.add_capture(0, metadata={
        SigMFFile.FREQUENCY_KEY: freq,
        SigMFFile.DATETIME_KEY: datetime.datetime.utcnow().isoformat()+'Z',
    })
    mete_fn = f'rx_{formatted_date}_{centerFreq}m.sigmf-meta'
    meta.tofile(str(Path(outDir).joinpath(mete_fn)))
    print("Closing the connection")
    writer.close()

async def datsrv():
    server = await asyncio.start_server(handle_client, '0.0.0.0', port)

    addr = server.sockets[0].getsockname()
    print(f'Serving on {addr}')

    return server


def set_freq_cmd():
    ret = []
    ret.append("rf-app -w 10d 0")
    freq = int(round(centerFreq*1e6))
    ret.append(f"rf-app -freq {freq} 5000000000")
    ret.append("rf-app -tg 8000 8000")
    ret.append("rf-app -w 10d 1")
    ret.append(f"rf-app -freq {freq} 5000000000")
    ret.append("rf-app -tg 8000 8000")
    return ret
    
async def shell(reader, writer):
    # dat_srv = await datsrv()    
    async def cmdline(cmd):
        writer.write(cmd)
        ret = await reader.readuntil(b"root@ai_prj:~# ")
        print(f"done: \n{ret.decode('ascii')}")
    def direct(cmd):
        writer.write(cmd)
    async def dump():
        ret = await reader.readuntil(b"root@ai_prj:~# ")
        print(f"done: \n{ret.decode('ascii')}")
        
    await reader.readuntil(b"login: ")
    writer.write("root\n")
    print("login ...")
    await reader.readuntil(b"Password: ")
    writer.write("root\n")
    await reader.readuntil(b"root@ai_prj:~# ")
    print("login")
    if init:
        await cmdline("cp /run/media/mmcblk0p1/*.txt .\n")
        await cmdline("cp /run/media/mmcblk0p1/run_all.sh .\n")
        await cmdline("sh run_all.sh\n")
    cmds = set_freq_cmd()
    for cmd in cmds:
        await cmdline(cmd+"\n")
    print("transmission start")
    direct(f"rf-app -rx_socket_dma 64 {IPaddr}\n")
    data_server()
    await dump()
    print("transmission complete ok")
    writer.close()

if __name__ == "__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:], "if:d:p:a:", ["init","freq=", "dir=", "port=", "ip="])
    except getopt.GetoptError as err:
        # print help information and exit:
        print(err)  
        usage()
        sys.exit(2)
    for opt, arg in opts:
        if opt in ('-f', '--freq'):
            centerFreq = float(arg)
        elif opt in ('-d', '--dir'):
            outDir = arg
        elif opt in ('-p', '--port'):
            port = int(arg)
        elif opt in ('-i', '--init'):
            init = True
        elif opt in ('-a', '--ip'):
            IPaddr = arg
            setip = True

    if not setip:
        if not get_ip():
            print("bad ip address")
            sys.exit(3)
            
    loop = asyncio.get_event_loop()
    
    coro = telnetlib3.open_connection('192.168.1.100', 23, shell=shell)
    
    reader, writer = loop.run_until_complete(coro)
    loop.run_until_complete(writer.protocol.waiter_closed)
    
    print("End")
    